from contextlib import contextmanager
from opentelemetry import trace
from opentelemetry.trace import INVALID_SPAN
from opentelemetry.propagate import extract, set_global_textmap, get_global_textmap
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace.sampling import TraceIdRatioBased, ParentBased
from pydantic import BaseModel, Field

from common.utils import logger

class TelemetryConf(BaseModel):
    name: str = Field(alias="Name")
    endpoint: str = Field(alias="Endpoint")
    sampler: float = Field(alias="Sampler")
    batcher: str = Field(alias="Batcher")

def init_telemetry(c: TelemetryConf):
    # 新增：标准化传播协议
    set_global_textmap(TraceContextTextMapPropagator())
    
    # 1. 资源定义（增强语义属性）
    resource = Resource.create(attributes={
        "service.name": c.name,
        "telemetry.sdk.language": "python",
        "service.version": "1.0.0"
    })

    # 2. 创建TracerProvider
    provider = TracerProvider(
        resource=resource,
        sampler=ParentBased(TraceIdRatioBased(c.sampler))
    )

    # 3. 替换为OTLP导出器（协议兼容性更优）
    otlp_exporter = OTLPSpanExporter(
        endpoint=c.endpoint,
        insecure=True if not c.endpoint.startswith("https") else False
    )

    # 4. 批量处理器（性能优化）
    provider.add_span_processor(BatchSpanProcessor(otlp_exporter))

    # 5. 全局设置
    trace.set_tracer_provider(provider)

class TracingUtil:
    def __init__(self):
        self.tracer = trace.get_tracer(__name__)
        self.propagator = get_global_textmap()
    
    @contextmanager
    def __call__(self, context):
        span = self._get_or_create_span(context)
        if span is None:
            logger.my_logger().error("Failed to get or create span: context metadata is invalid or missing")
            yield None
            return
            
        with trace.use_span(span, end_on_exit=True):
            yield span
    
    def _get_or_create_span(self, context):
        try:
            metadata = dict(context.invocation_metadata())
            ctx = extract(metadata) if metadata else None
            return self.tracer.start_span(
                "auto-generated-span",
                context=ctx or trace.Context()
            )
        except Exception as e:
            logger.my_logger().error(f"Span extraction failed: {str(e)}")
            return None
    
    def inject_trace_metadata(self, metadata: dict):
        """将当前trace信息注入到metadata字典中"""
        current_span = trace.get_current_span()
        if current_span is INVALID_SPAN:
            logger.my_logger().error("Cannot inject trace metadata: no active span found")
            return metadata
            
        try:
            carrier = {}
            self.propagator.inject(carrier)
            metadata.update(carrier)
        except Exception as e:
            logger.my_logger().error(f"Failed to inject trace metadata: {str(e)}")
        return metadata